"""多模型融合推理管道"""
from typing import Dict, Any, List
from fin_senti_entity_platform.utils.config_loader import config_loader
from fin_senti_entity_platform.utils.logger import get_logger
from fin_senti_entity_platform.model_development.sentiment_analysis.deepseek_sentiment import DeepSeekSentimentModel
from fin_senti_entity_platform.model_development.sentiment_analysis.rnn_trainer import RNNModel, RNNModelTrainer
from fin_senti_entity_platform.model_development.entity_recognition.bert_crf_trainer import BertCRFTrainer

logger = get_logger(__name__)

class FusionPipeline:
    """多模型融合推理管道，整合多种模型的预测结果"""
    
    def __init__(self):
        """初始化融合推理管道"""
        self.models = {}
        self.weights = {}
        self._load_models()
        self._set_model_weights()
    
    def _load_models(self):
        """加载所有需要的模型"""
        logger.info("开始加载融合推理所需的模型...")
        
        # 加载DeepSeek-7B大语言模型
        try:
            self.models['deepseek'] = DeepSeekSentimentModel()
            logger.info("成功加载DeepSeekSentimentModel")
        except Exception as e:
            logger.error(f"加载DeepSeekSentimentModel失败: {e}")
        
        # 加载RNN时序模型
        try:
            rnn_trainer = RNNModelTrainer()
            rnn_model_path = f"{config_loader.config['project_paths']['model_saves']}/rnn_model"
            rnn_trainer.load_model(rnn_model_path)
            self.models['rnn'] = rnn_trainer
            logger.info("成功加载RNNModel")
        except Exception as e:
            logger.error(f"加载RNNModel失败: {e}")
        
        # 加载实体识别模型
        try:
            entity_trainer = BertCRFTrainer()
            self.models['entity'] = entity_trainer
            logger.info("成功加载BertCRFTrainer")
        except Exception as e:
            logger.error(f"加载BertCRFTrainer失败: {e}")
    
    def _set_model_weights(self):
        """设置模型权重"""
        # 可以根据模型性能动态调整权重
        self.weights = {
            'deepseek': config_loader.config['model']['fusion']['weights'].get('deepseek', 0.5),
            'rnn': config_loader.config['model']['fusion']['weights'].get('rnn', 0.3),
            # 其他模型的权重...
        }
        
        # 归一化权重
        total_weight = sum(self.weights.values())
        for model_name in self.weights:
            self.weights[model_name] /= total_weight
    
    def recognize_entities(self, text: str) -> List[Dict[str, Any]]:
        """识别文本中的实体"""
        if 'entity' in self.models:
            try:
                entities = self.models['entity'].predict(text)
                return entities
            except Exception as e:
                logger.error(f"实体识别失败: {e}")
                return []
        return []
    
    def _get_individual_predictions(self, text: str, entities: List[Dict[str, Any]]) -> Dict[str, Dict[str, float]]:
        """获取各个模型的单独预测结果"""
        predictions = {}
        
        # DeepSeek模型预测
        if 'deepseek' in self.models:
            try:
                pred = self.models['deepseek'].predict(text, entities)
                predictions['deepseek'] = pred['probabilities']
            except Exception as e:
                logger.error(f"DeepSeek模型预测失败: {e}")
        
        # RNN模型预测
        if 'rnn' in self.models:
            try:
                pred = self.models['rnn'].predict(text, entities)
                predictions['rnn'] = pred['probabilities']
            except Exception as e:
                logger.error(f"RNN模型预测失败: {e}")
        
        return predictions
    
    def _fuse_predictions(self, predictions: Dict[str, Dict[str, float]]) -> Dict[str, float]:
        """融合多个模型的预测结果"""
        if not predictions:
            return {label: 1.0/len(config_loader.config['model']['sentiment_labels']) for label in config_loader.config['model']['sentiment_labels']}

        fused_probabilities = {label: 0.0 for label in config_loader.config['model']['sentiment_labels']}
        
        # 加权平均
        for model_name, model_pred in predictions.items():
            weight = self.weights.get(model_name, 0.0)
            if weight > 0:
                for label in config_loader.config['model']['sentiment_labels']:
                    if label in model_pred:
                        fused_probabilities[label] += weight * model_pred[label]
        
        # 归一化
        total = sum(fused_probabilities.values())
        if total > 0:
            for label in fused_probabilities:
                fused_probabilities[label] /= total
        
        return fused_probabilities
    
    def predict(self, text: str, entities: List[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        融合预测文本情感
        
        Args:
            text: 输入文本
            entities: 可选的预识别实体列表
            
        Returns:
            预测结果，包含情感标签、置信度和各模型的预测概率
        """
        # 如果没有提供实体，自动识别
        if entities is None or not entities:
            entities = self.recognize_entities(text)
        
        # 获取各个模型的预测结果
        individual_predictions = self._get_individual_predictions(text, entities)
        
        # 融合预测结果
        fused_probabilities = self._fuse_predictions(individual_predictions)
        
        # 确定最终的情感标签
        sentiment_labels = config_loader.config['model']['sentiment_labels']
        predicted_label = max(fused_probabilities.items(), key=lambda x: x[1])[0]
        confidence = fused_probabilities[predicted_label]
        
        # 返回结果
        result = {
            'text': text,
            'entities': entities,
            'sentiment': predicted_label,
            'confidence': confidence,
            'probabilities': fused_probabilities,
            'individual_predictions': individual_predictions
        }
        
        return result
    
    def batch_predict(self, texts: List[str], batch_size: int = 8) -> List[Dict[str, Any]]:
        """
        批量预测文本情感
        
        Args:
            texts: 文本列表
            batch_size: 批次大小
            
        Returns:
            预测结果列表
        """
        results = []
        
        # 分批处理
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            batch_results = []
            
            # 对每批文本进行预测
            for text in batch_texts:
                try:
                    result = self.predict(text)
                    batch_results.append(result)
                except Exception as e:
                    logger.error(f"预测文本 '{text[:20]}...' 失败: {e}")
                    # 添加默认结果
                    batch_results.append({
                        'text': text,
                        'entities': [],
                        'sentiment': '中性',
                        'confidence': 0.5,
                        'probabilities': {label: 1.0/len(config_loader.config['model']['sentiment_labels']) for label in config_loader.config['model']['sentiment_labels']},
                        'individual_predictions': {}
                    })
            
            results.extend(batch_results)
        
        return results
    
    def update_model_weights(self, new_weights: Dict[str, float]):
        """
        更新模型权重
        
        Args:
            new_weights: 新的模型权重字典
        """
        for model_name, weight in new_weights.items():
            if model_name in self.weights:
                self.weights[model_name] = weight
        
        # 重新归一化
        self._set_model_weights()
        logger.info(f"更新后的模型权重: {self.weights}")
    
    def warm_up(self, sample_texts: List[str] = None):
        """
        预热模型以提高首次推理速度
        
        Args:
            sample_texts: 用于预热的样本文本列表
        """
        if sample_texts is None:
            sample_texts = ["这是一条测试文本，用于预热模型。"] * 2
        
        logger.info("开始预热模型...")
        for text in sample_texts:
            try:
                self.predict(text)
            except Exception as e:
                logger.warning(f"模型预热时出错: {e}")
        logger.info("模型预热完成")